Skip to content

Support shuffle control in JAX random choice#3238

Merged
FlorianPfaff merged 2 commits into
mainfrom
fix-jax-choice-shuffle-20260627
Jun 27, 2026
Merged

Support shuffle control in JAX random choice#3238
FlorianPfaff merged 2 commits into
mainfrom
fix-jax-choice-shuffle-20260627

Conversation

@FlorianPfaff

@FlorianPfaff FlorianPfaff commented Jun 27, 2026

Copy link
Copy Markdown
Owner

Summary

  • add the NumPy/PyTorch-compatible shuffle flag to the JAX random choice facade
  • sample indices explicitly so shuffle=False with replace=False preserves population and axis order
  • add focused regression tests for order preservation and invalid shuffle controls

Bug fixed

pyrecest._backend.jax.random.choice(..., shuffle=False) was previously forwarded to jax.random.choice, where shuffle is not part of PyRecEst's facade-normalized behavior. That made code portable across the NumPy/PyTorch backends fail or diverge under JAX.

Validation

  • Locally exercised the revised JAX choice logic for one-dimensional and axis=1 array populations, including invalid shuffle controls.
  • Not run as a full repository pytest suite here; the execution sandbox cannot resolve github.com, so I could not clone/install the complete repository in this environment.

@FlorianPfaff FlorianPfaff enabled auto-merge (squash) June 27, 2026 21:01
@github-actions

github-actions Bot commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

✅MegaLinter analysis: Success

Descriptor Linter Files Fixed Errors Warnings Elapsed time
✅ COPYPASTE jscpd yes no no 70.02s
✅ JSON prettier 7 0 0 0 1.38s
✅ JSON v8r 7 0 0 5.31s
✅ MARKDOWN markdownlint 68 0 0 0 1.65s
✅ MARKDOWN markdown-table-formatter 68 0 0 0 0.87s
✅ PYTHON black 1051 50 0 0 68.66s
✅ PYTHON isort 1051 72 0 0 4.18s
✅ REPOSITORY checkov yes no no 51.36s
✅ REPOSITORY gitleaks yes no no 5.78s
✅ REPOSITORY git_diff yes no no 0.16s
✅ REPOSITORY secretlint yes no no 29.9s
✅ REPOSITORY syft yes no no 2.79s
✅ REPOSITORY trivy-sbom yes no no 7.18s
✅ REPOSITORY trufflehog yes no no 25.51s
✅ YAML prettier 11 0 0 0 0.74s
✅ YAML v8r 11 0 0 12.78s
✅ YAML yamllint 11 0 0 0.56s

Notices

📣 MegaLinter 9.5.0 is out! Discover the new features and security recommendations in the release announcement. (Skip this info by defining SECURITY_SUGGESTIONS: false)

See detailed reports in MegaLinter artifacts

Your project could benefit from a custom flavor, which would allow you to run only the linters you need, and thus improve runtime performances. (Skip this info by defining FLAVOR_SUGGESTIONS: false)

  • Documentation: Custom Flavors
  • Command: npx mega-linter-runner@9.5.0 --custom-flavor-setup --custom-flavor-linters PYTHON_BLACK,PYTHON_ISORT,COPYPASTE_JSCPD,JSON_V8R,JSON_PRETTIER,MARKDOWN_MARKDOWNLINT,MARKDOWN_MARKDOWN_TABLE_FORMATTER,REPOSITORY_CHECKOV,REPOSITORY_GIT_DIFF,REPOSITORY_GITLEAKS,REPOSITORY_SECRETLINT,REPOSITORY_SYFT,REPOSITORY_TRIVY_SBOM,REPOSITORY_TRUFFLEHOG,YAML_PRETTIER,YAML_YAMLLINT,YAML_V8R

MegaLinter is graciously provided by OX Security
Show us your support by starring ⭐ the repository

@FlorianPfaff FlorianPfaff merged commit 7c4f451 into main Jun 27, 2026
25 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@FlorianPfaff